package mcp

import (
	"bufio"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"os"
)

// Tool 定义MCP工具
type Tool struct {
	Name        string                 `json:"name"`
	Description string                 `json:"description"`
	InputSchema map[string]interface{} `json:"inputSchema"`
}

// ToolHandler 工具处理函数类型
type ToolHandler func(ctx context.Context, args map[string]interface{}) (interface{}, error)

// Server MCP服务器
type Server struct {
	name    string
	version string
	tools   map[string]Tool
	handlers map[string]ToolHandler
}

// NewServer 创建新的MCP服务器
func NewServer(name, version string) *Server {
	return &Server{
		name:     name,
		version:  version,
		tools:    make(map[string]Tool),
		handlers: make(map[string]ToolHandler),
	}
}

// AddTool 添加工具
func (s *Server) AddTool(tool Tool, handler ToolHandler) {
	s.tools[tool.Name] = tool
	s.handlers[tool.Name] = handler
}

// MCPRequest MCP请求结构
type MCPRequest struct {
	JSONRPC string                 `json:"jsonrpc"`
	ID      interface{}            `json:"id"`
	Method  string                 `json:"method"`
	Params  map[string]interface{} `json:"params,omitempty"`
}

// MCPResponse MCP响应结构
type MCPResponse struct {
	JSONRPC string      `json:"jsonrpc"`
	ID      interface{} `json:"id"`
	Result  interface{} `json:"result,omitempty"`
	Error   *MCPError   `json:"error,omitempty"`
}

// MCPError MCP错误结构
type MCPError struct {
	Code    int    `json:"code"`
	Message string `json:"message"`
}

// Run 运行MCP服务器
func (s *Server) Run() error {
	reader := bufio.NewReader(os.Stdin)
	
	for {
		line, err := reader.ReadString('\n')
		if err != nil {
			if err == io.EOF {
				break
			}
			return err
		}

		var req MCPRequest
		if err := json.Unmarshal([]byte(line), &req); err != nil {
			log.Printf("解析请求失败: %v", err)
			continue
		}

		response := s.handleRequest(req)
		
		responseData, err := json.Marshal(response)
		if err != nil {
			log.Printf("序列化响应失败: %v", err)
			continue
		}

		fmt.Println(string(responseData))
	}

	return nil
}

// handleRequest 处理MCP请求
func (s *Server) handleRequest(req MCPRequest) MCPResponse {
	switch req.Method {
	case "initialize":
		return MCPResponse{
			JSONRPC: "2.0",
			ID:      req.ID,
			Result: map[string]interface{}{
				"protocolVersion": "2024-11-05",
				"capabilities": map[string]interface{}{
					"tools": map[string]interface{}{},
				},
				"serverInfo": map[string]interface{}{
					"name":    s.name,
					"version": s.version,
				},
			},
		}

	case "tools/list":
		tools := make([]Tool, 0, len(s.tools))
		for _, tool := range s.tools {
			tools = append(tools, tool)
		}
		return MCPResponse{
			JSONRPC: "2.0",
			ID:      req.ID,
			Result: map[string]interface{}{
				"tools": tools,
			},
		}

	case "tools/call":
		toolName, ok := req.Params["name"].(string)
		if !ok {
			return MCPResponse{
				JSONRPC: "2.0",
				ID:      req.ID,
				Error: &MCPError{
					Code:    -32602,
					Message: "工具名称无效",
				},
			}
		}

		handler, exists := s.handlers[toolName]
		if !exists {
			return MCPResponse{
				JSONRPC: "2.0",
				ID:      req.ID,
				Error: &MCPError{
					Code:    -32601,
					Message: fmt.Sprintf("工具不存在: %s", toolName),
				},
			}
		}

		args, ok := req.Params["arguments"].(map[string]interface{})
		if !ok {
			args = make(map[string]interface{})
		}

		result, err := handler(context.Background(), args)
		if err != nil {
			return MCPResponse{
				JSONRPC: "2.0",
				ID:      req.ID,
				Error: &MCPError{
					Code:    -32603,
					Message: err.Error(),
				},
			}
		}

		return MCPResponse{
			JSONRPC: "2.0",
			ID:      req.ID,
			Result: map[string]interface{}{
				"content": []map[string]interface{}{
					{
						"type": "text",
						"text": result,
					},
				},
			},
		}

	default:
		return MCPResponse{
			JSONRPC: "2.0",
			ID:      req.ID,
			Error: &MCPError{
				Code:    -32601,
				Message: fmt.Sprintf("不支持的方法: %s", req.Method),
			},
		}
	}
}
